# 抽象的DFT矩阵分解基类
import numpy as np
from solutions.digital_fourier_transform_base import DigitalFourierTransformerBase
import pickle
import os


class DFT_1(DigitalFourierTransformerBase):
    @property
    def __name__(self):
        return "Problem_1"

    def identify(self, A_k: np.ndarray, rate: float = 0.00001) -> np.ndarray:
        """
        使用一波最原始的方法，取出每行的绝对值最大的元素，将其放大某个倍数，然后将第二大的也放大这个位数，将其他的缩小某个倍数。
        """
        res = A_k - self.max_two(A_k)
        return res * (1 - rate) + self.max_two(A_k) * (1 + rate)

    def grad(self, k: int, m: float = 3):
        return (
            -(self.beta * self.dft - self._matmul(self._A_K, self.K))
            @ np.conjugate(self._matmul(self._A_K, k))
            - self._matmul(self._A_K, k)
            @ np.conjugate((self.beta * self.dft - self._matmul(self._A_K, self.K)))
            * m
        )

    def grad_beta(self):
        return np.real(
            np.trace(
                (self.beta * self.dft - self._matmul(self._A_K, self.K))
                @ np.conjugate(self.dft)
            )
        )

    def train(self, lr: float = 0.01, epoch: int = 1000, store_path: os.PathLike = './result', **kwargs):
        error = 100000
        error_list = []
        beta_list = []
        A_K_list = []
        for i in range(epoch):
            try:
                A_K_temp = [
                    self.identify(x - lr * self.grad(k))
                    for k, x in enumerate(self._A_K)
                ]
                beta_temp = self.beta - lr * self.grad_beta()
                if np.abs(self.error - error) < 1e-8:
                    break

                if i % 100 == 0:
                    # error = self.error
                    print(f"epoch:{i}, error:{self.error}")
                error = self.error
                self._A_K = A_K_temp
                self.beta = beta_temp
                error_list.append(error)
                A_K_list.append(A_K_temp)
                beta_list.append(beta_temp)
            except:
                raise Exception("TrainError")
        self._A_K = [self.max_two(A_K_i) for A_K_i in self._A_K]
        pickle.dump((self, error_list, A_K_list, beta_list),
                    open(os.path.join(store_path, f"K{self.K}_N{self.N}_lr{lr}_"+self.__name__+'.pkl'), 'wb'))
